import os.path

import cv2
from ultralytics import YOLO, solutions
import tkinter as tk
from tkinter import filedialog
import os

# 定义参数
# video = 'video/4.mp4'   # 固定地址
# video_result = 'result/video_result.mp4'   # 固定地址
huatu = 'result/diyizhen.jpg'
WN = 'Select Line Points'   #窗口名称
model = YOLO('yolo12n.pt')
mubiao = [0]   # 需检测的目标，这里仅仅检测人
dian = []   # 划线要的两个点坐标

# 获取视频文件路径
root = tk.Tk()
root.withdraw()
video = filedialog.askopenfilename(
    title="选择输入视频",
    filetypes=[("视频文件", "*.mp4 *.avi *.mkv")]
)

video_result = filedialog.askdirectory(
    title="选择输出目录"
)
video_result = os.path.join(video_result,'video_result.mp4')

v1 = cv2.VideoCapture(video)  # 读取视频
# 视频报错检测
# if not v1.isOpened():
#     print(f"错误：无法打开视频文件 {video}")
#     exit()
assert v1.isOpened(), '无法打开视频文件'

# print(v1.get(3),v1.get(4),v1.get(5))   # 获取视频的宽度、高度和帧率
kuan = v1.get(3)
gao = v1.get(4)
FPS = v1.get(5)
org = (int(kuan/10), int(gao/10))   # 标题坐标位置
# print(org)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')   # 保存格式
out = cv2.VideoWriter(
    video_result,
    fourcc,
    FPS,
    (int(kuan),int(gao))
)   # 定义视频保存方式

def zuobiao(event,x,y,flags,para):   # 回调函数
    if event == cv2.EVENT_LBUTTONDOWN:   # 按下左键后
        dian.append((x,y))   # 获取左键时鼠标x，y坐标
        if len(dian) > 2:   # 当坐标大于两个时，删除第一个坐标
            dian.pop(0)
        if len(dian) == 2:
            return dian

# 获取第一帧，用于获取坐标
ret1, frame1 = v1.read()
if not ret1:
    print('没有获取到视频')
    exit()
cv2.imwrite(huatu,frame1)
one = cv2.imread(huatu)
cv2.putText(one,'请在要检测的地方点两下以画成直线',org,cv2.FONT_HERSHEY_SIMPLEX,3,(0,0,255),2)
# 下面三个的窗口名要一致才能在图片中获取坐标
cv2.namedWindow(WN, cv2.WINDOW_FREERATIO)
cv2.setMouseCallback(WN, zuobiao)
cv2.imshow(WN,one)
cv2.waitKey(0)
cv2.destroyAllWindows()

# yolo对象计数函数
counter = solutions.ObjectCounter(
    show=True,                    # 【显示输出】控制是否在界面显示实时计数结果（默认True）
    region=dian,          # 【指定检测区域】传入定义的检测区域坐标（例如入口/出口的矩形坐标）
    model="yolo12n.pt",           # 【模型路径】指定使用的YOLO检测模型（支持.ypt/.pt格式）
    classes=[0],              # 【可选】限定计数类别（例如COCO数据集中0=人，2=车）
    tracker="bytetrack.yaml",      # 【可选】指定跟踪器配置文件（如bytetrack.yaml实现高性能跟踪,默认追踪器为botsort.yaml）
    conf = 0.1,               # 置信度设置
)


while True:
    ret, frame = v1.read()  # 获取视频每一帧

    if not ret:  # 读取帧失败（视频结束），退出循环
        break

    # genzong = model.track(frame, persist=True, classes = mubiao, device = '0')   # track是用于实例跟踪的方法，能持续对目标进行跟踪并生成带有唯一的ID；其中的persist参数，开启后能连续追踪目标
    # v2 = genzong[0].plot()   # track回来的并不是图片，而是一串数据：检测结果，而plot是在图像上绘制检测结果
    cv2.namedWindow('Ultralytics Solutions',cv2.WINDOW_FREERATIO)
    v3 = counter(frame).plot_im   # 执行计数并且获得结果  报错修复：添加plot_im,显示处理后的帧，输出数据类型为numpy.ndarray

    # cv2.imshow('shiping', v3)   # imshow需要接收numpy.ndarray数据类型
    out.write(frame)   # 将处理后的帧写入输出视频
    if cv2.waitKey(1) == ord("q"):  # 按q结束运行，waitKey值为0时，程序在按下指定按键前无法运行其他内容，waitKey为1时，代表程序每0.1秒检测一次是否按下指定按键
        break

cv2.destroyAllWindows()  # 关闭所有opencv窗口
v1.release()
out.release()
